#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sys/time.h>

void Conv14x4(float *dst, float *input, float *weight, size_t src_depth, size_t dst_step, size_t dst_depth, size_t width);

void Conv8x8(float *input, float *weight, float *output, size_t step, size_t offset);

void main()
{
    int iw = 448;
    int ih = 448;
    int ic = 64;
    int oc = 32;
    int kw = 3;
    int kh = 3;
    int ow = 448;
    int oh = 448;

    int packed_input_size = ow * oh * kw * kh * ic;
    float *packed_input_1 = (float *) malloc(packed_input_size * sizeof(float));
    float *packed_input_2 = (float *) malloc(packed_input_size * sizeof(float));
    for (int i = 0; i < packed_input_size; ++i) {
        packed_input_1[i] = (float) (i % 10);
        packed_input_2[i] = packed_input_1[i];
    }

    int packed_weight_size = oc * kw * kh * ic;
    float *packed_weight_1 = (float *) malloc(packed_weight_size * sizeof(float));
    float *packed_weight_2 = (float *) malloc(packed_weight_size * sizeof(float));
    for (int i = 0; i < packed_weight_size; ++i) {
        packed_weight_1[i] = (float) (i % 5 + 2);
        packed_weight_2[i] =packed_weight_1[i];
    }

    int out_size = ow * oh * oc;
    float *output_1 = (float *) malloc(out_size * sizeof(float));
    float *output_2 = (float *) malloc(out_size * sizeof(float));
    memset(output_1, 0, out_size * sizeof(float));
    memset(output_2, 0, out_size * sizeof(float));

    // 14x4
    int batches = ow * oh / 14;
    struct timeval start, end;
    double duration = 0;
    gettimeofday(&start, NULL);
    for (int i = 0; i < batches; ++i) {
        Conv14x4(output_1, packed_input_1 + i * kh * kw * 4, packed_weight_1, kw * kh * ic / 4, ow * oh * 4, oc / 4, 0);
    }
    gettimeofday(&end, NULL);
    duration = (end.tv_sec - start.tv_sec) * 1000000 + (end.tv_usec - start.tv_usec);
    printf("%f\n", duration);

    // 8x8
    batches = ow * oh * oc / 8 / 8;
    int ksize = kw * kh;
    int step = ic / 4;
    gettimeofday(&start, NULL);
    for (int i = 0; i < batches; ++i) {
        Conv8x8(packed_input_2 + i * kw * kh * 8, packed_weight_2 + i % (ic / 8), output_2, step * ksize, oc * 4);
    }
    gettimeofday(&end, NULL);
    duration = (end.tv_sec - start.tv_sec) * 1000000 + (end.tv_usec - start.tv_usec);
    printf("%f\n", duration);
}